package top.jfunc.common.db.data;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.sql.ParameterMetaData;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.function.Consumer;
import java.util.function.IntConsumer;

import javax.sql.DataSource;

/**
 * 分页数据采集
 *
 * @author <a href="https://github.com/TAKETODAY">Harry Yang</a>
 * @since 1.0 2022/11/25 11:03
 */
final class PageableSQLDataExtractor extends DataExtractor {
  private static final Logger logger = LoggerFactory.getLogger(PageableSQLDataExtractor.class);

  private static final int DEFAULT_PAGE_SIZE = 5000;

  private final String firstPageSql;
  private final String remainingPagesSql;

  private final int pageSize;
  private final DataSource dataSource;
  private final List<String> sortKeys;

  private final IntConsumer pageFinished;

  PageableSQLDataExtractor(int pageSize,
      PagingQueryProvider queryProvider,
      DataSource dataSource,
      List<String> sortKeys,
      Consumer<Throwable> exceptionConsumer,
      PreparedStatementSetter statementSetter,
      IntConsumer successCallback, IntConsumer pageFinished) {
    super(dataSource, exceptionConsumer, statementSetter, successCallback);
    this.pageSize = pageSize;
    this.sortKeys = sortKeys;
    this.dataSource = dataSource;
    this.pageFinished = pageFinished;
    this.firstPageSql = queryProvider.generateFirstPageQuery(pageSize);
    this.remainingPagesSql = queryProvider.generateRemainingPagesQuery(pageSize);
  }

  /**
   * 该方法 没有数据会返回
   *
   * @param handler 数据行回调处理
   */
  @Override
  public int processRow(RowCallbackHandler handler) {
    if (logger.isDebugEnabled()) {
      logger.debug("SQL used for reading first page: [{}]", firstPageSql);
    }
    PageRowHandler rowHandler = new PageRowHandler(handler, statementSetter, pageSize, sortKeys);
    DataExtractor.fetch(firstPageSql, dataSource)
        .config(statementSetter)
        .errorHandling(exceptionConsumer)
        .subscribe(rowHandler);

    // 判断是否只有一页的数据
    if (rowHandler.isFinishedReading()) {
      // 只有一页数据
      pageFinished(1);
      return rowHandler.processedRows;
    }

    // 超过一页的数据
    int processedRows = rowHandler.processedRows;

    int current = 2;

    DataExtractor remaining =
        DataExtractor.fetch(remainingPagesSql, dataSource)
            .errorHandling(exceptionConsumer)
            .config(rowHandler)
            .build();
    do {
      if (logger.isDebugEnabled()) {
        logger.debug("Reading page {}", current);
      }
      rowHandler.reset();
      remaining.processRow(rowHandler);

      if (rowHandler.processedRows != 0) {
        pageFinished(current);
        // 不是刚好读完，就读了一页，反之上一次就是最后一页
        current++;
        processedRows += rowHandler.processedRows;
      }
    }
    while (rowHandler.hasRemaining());

    if (successCallback != null) {
      successCallback.accept(processedRows);
    }
    return processedRows;
  }

  private void pageFinished(int current) {
    if (pageFinished != null) {
      pageFinished.accept(current);
    }
  }

  public static Builder forBuilder() {
    return new Builder();
  }

  // Helper

  static final class PageRowHandler implements RowCallbackHandler, PreparedStatementSetter {
    private final RowCallbackHandler handler;
    private final PreparedStatementSetter statementSetter;
    private final int pageSize;
    // 从 报表库读取的行数
    private int processedRows = 0;
    private final List<String> sortKeys;
    private final ArrayList<Object> sortValues;

    /**
     * 当前分页读取完毕标志
     */
    private boolean pageFinished;

    public PageRowHandler(RowCallbackHandler handler,
        PreparedStatementSetter statementSetter, int pageSize, List<String> sortKeys) {
      this.handler = handler;
      this.statementSetter = statementSetter;
      this.pageSize = pageSize;
      this.sortKeys = sortKeys;
      this.sortValues = new ArrayList<>(sortKeys.size());
    }

    @Override
    public void setValues(PreparedStatement ps) throws SQLException {
      if (statementSetter != null) {
        statementSetter.setValues(ps);
      }

      if (sortValues.isEmpty()) {
        return;
      }

      ParameterMetaData parameterMetaData = ps.getParameterMetaData();
      int parameterCount = parameterMetaData.getParameterCount();
      int idx = parameterCount - sortKeys.size() + 1;
      for (Object sortValue : sortValues) {
        ps.setObject(idx++, sortValue);
      }
    }

    @Override
    public void finishedReading(int rows) throws Exception {
      handler.finishedReading(rows);
      this.processedRows = rows;
      this.pageFinished = true;
    }

    @Override
    public void processRow(ResultSet rs, int rowNumber) throws Exception {
      handler.processRow(rs, rowNumber);

      if (rowNumber == pageSize) {
        // 最后一行 记录下一次分页位置
        sortValues.clear();
        for (String sortKey : sortKeys) {
          Object sortValue = rs.getObject(sortKey);
          sortValues.add(sortValue);
        }
      }
      // else 少于分页大小的不需要记录，因为已经读取完毕，没有更多的页数了
    }

    /**
     * 判断是不是全部数据已经读取完毕
     */
    public boolean isFinishedReading() {
      // 读取的行数已经小于分页数了，说明所有数据已经读取完毕
      return pageFinished && processedRows < pageSize;
    }

    public boolean hasRemaining() {
      // 判断是否只有一页的数据
      return processedRows == pageSize;
    }

    public void reset() {
      processedRows = 0;
      this.pageFinished = false;
    }

  }

  // Builder
  public static class Builder extends DataExtractorBuilderSupport<Builder> implements PageableBuilder {
    private int pageSize = DEFAULT_PAGE_SIZE;
    private IntConsumer pageFinished;
    private final PagingQueryProvider queryProvider = new MySqlPagingQueryProvider();

    @Override
    protected Builder self() {
      return this;
    }

    public Builder pageSize(int pageSize) {
      this.pageSize = pageSize;
      return this;
    }

    public Builder select(String selectClause) {
      queryProvider.setSelectClause(selectClause);
      return this;
    }

    public Builder from(String fromClause) {
      queryProvider.setFromClause(fromClause);
      return this;
    }

    public Builder where(String whereClause) {
      queryProvider.setWhereClause(whereClause);
      return this;
    }

    public Builder groupBy(String groupClause) {
      queryProvider.setGroupClause(groupClause);
      return this;
    }

    /**
     * 排序
     */
    public Builder orderBy(Map<String, Order> sortKeys) {
      queryProvider.setSortKeys(sortKeys);
      return this;
    }

    /**
     * 排序
     *
     * @param sortKey 排序的字段
     * @param order 排序的方向
     */
    public Builder orderBy(String sortKey, Order order) {
      return orderBy(Collections.singletonMap(sortKey, order));
    }

    /**
     * 排序，默认升序
     *
     * @param sortKey 排序的字段
     */
    public Builder orderBy(String sortKey) {
      return orderBy(sortKey, Order.ASCENDING);
    }

    @Override
    public PageableBuilder onPageFinished(IntConsumer current) {
      this.pageFinished = current;
      return this;
    }

    @Override
    public PageableSQLDataExtractor build() {
      List<String> sortKeys = new ArrayList<>(queryProvider.getSortKeys().keySet());
      return new PageableSQLDataExtractor(pageSize, queryProvider,
          dataSource, sortKeys, exceptionConsumer,
          statementPostProcessor, successCallback, pageFinished);
    }

  }

}
